from ultralytics import YOLO
import cv2
import os
import numpy as np
from pathlib import Path

# 加载预训练的YOLOv11n模型
model = YOLO(r"D:\songlin\Python\PythonTraining3\ultralytics\runs\segment\train\weights\best.pt")

# 创建保存裁剪结果的目录
save_dir = r"D:\songlin\Python\PythonTraining3\ultralytics\segmented_crops"
os.makedirs(save_dir, exist_ok=True)

def get_min_bounding_box(image_path):
    image_name = Path(image_path).name
    # 对指定的图像进行推理
    results = model.predict(
        source=image_path,
        conf=0.45,
        iou=0.6,
        imgsz=640,
        save=False,  # 关闭自动保存，我们将手动处理[8](@ref)
        show=False
    )

    # 处理每个检测结果
    for i, result in enumerate(results):
        # 读取原始图像
        img = cv2.imread(result.path)
        img_height, img_width = img.shape[:2]  # 获取图像尺寸用于边界检查[4](@ref)

        # 获取分割掩码
        masks = result.masks

        if masks is not None:
            for j, mask in enumerate(masks):
                # 获取分割多边形坐标
                mask_data = mask.xy[0].astype(np.int32)

                # 计算最小外接矩形
                rect = cv2.minAreaRect(mask_data)
                box_points = cv2.boxPoints(rect).astype(np.int32)

                # 计算裁剪区域边界
                x_coords = box_points[:, 0]
                y_coords = box_points[:, 1]
                x1, x2 = max(0, min(x_coords)), min(img_width, max(x_coords))  # 边界裁剪[4](@ref)
                y1, y2 = max(0, min(y_coords)), min(img_height, max(y_coords))  # 边界裁剪[4](@ref)

                # 创建图像副本用于绘制
                debug_img = img.copy()

                # 绘制最小外接矩形（绿色）
                cv2.drawContours(debug_img, [box_points], 0, (0, 255, 0), 2)

                # 绘制裁剪矩形（红色）
                cv2.rectangle(debug_img, (x1, y1), (x2, y2), (0, 0, 255), 2)

                # 保存调试图像
                debug_save_dir = os.path.join(save_dir, "debug_image")
                os.makedirs(debug_save_dir, exist_ok=True)
                debug_path = os.path.join(debug_save_dir, f"debug_{image_name}_{i}_{j}.jpg")
                cv2.imwrite(debug_path, debug_img)
                print(f"保存调试图像至: {debug_path}")

                # # 裁剪并保存分割区域
                # if x2 > x1 and y2 > y1:  # 确保是有效区域
                #     crop = img[y1:y2, x1:x2]
                #     crop_path = os.path.join(save_dir, f"crop_{i}_{j}.jpg")
                #     cv2.imwrite(crop_path, crop)
                #     print(f"保存裁剪图像至: {crop_path}")
                # else:
                #     print(f"警告: 跳过无效裁剪区域 {i}_{j}")

                # 裁剪并保存分割区域
                if x2 > x1 and y2 > y1:  # 确保是有效区域
                    # 获取最小外接矩形的四个顶点坐标
                    src_points = box_points.astype(np.float32)

                    # 计算目标矩形的宽度和高度
                    width = int(np.linalg.norm(src_points[0] - src_points[1]))
                    height = int(np.linalg.norm(src_points[1] - src_points[2]))

                    # 定义目标矩形的四个顶点
                    dst_points = np.array([
                        [0, 0],
                        [width - 1, 0],
                        [width - 1, height - 1],
                        [0, height - 1]
                    ], dtype=np.float32)

                    # 计算透视变换矩阵
                    M = cv2.getPerspectiveTransform(src_points, dst_points)

                    # 应用透视变换进行裁剪
                    crop = cv2.warpPerspective(img, M, (width, height))

                    # 如果高度大于宽度，则旋转90度
                    if crop.shape[0] > crop.shape[1]:
                        crop = cv2.rotate(crop, cv2.ROTATE_90_CLOCKWISE)

                    crop_path = os.path.join(save_dir, f"crop_{image_name}_{i}_{j}.jpg")
                    cv2.imwrite(crop_path, crop)
                    print(f"保存裁剪图像至: {crop_path}")
                else:
                    print(f"警告: 跳过无效裁剪区域 {i}_{j}")
        else:
            print("未检测到分割掩码")

    print(f"{image_name}截取后图片已保存至: {save_dir}")

image_folder = r"D:\songlin\data\8.图像分割相关\电子游标卡尺\电子游标卡尺dataset\images"
for root, dirs, files in os.walk(image_folder):
    for file in files:
        if file.endswith((".jpg", ".jpeg", ".png")):
            image_path = os.path.join(root, file)
            get_min_bounding_box(image_path)
